import matplotlib.pyplot as plt
import numpy as np

class show():
    def __init__(self, legend, xlabel=None, ylabel=None, xlim=None, ylim=None) -> None:
        self.xlabel = xlabel
        self.ylabel = ylabel
        self.xlim = xlim
        self.ylim = ylim
        self.legend = legend
        self.data = np.zeros(shape = (1,len(legend))) # 存储每个回合的数据，列是时间，行是不同的数据
        # self.data = [[0]*len(legend)]
        self.color = ["black", "red", "blue", "green", "orange", "yellow", "purples"]

    def add(self, args):
        if len(args)!=len(self.legend):
            assert(f"lenth of args {len(args)} no match length of legend{len(args)}")
        data:list = self.data.tolist()
        data.append(args)
        # print(data)
        self.data = np.array(data)
    
    def show_train(self):
        plt.cla()
        if self.xlim!=None:
            plt.xlim(self.xlim)
        if self.ylim!=None:
            plt.ylim(self.ylim)
        plt.xlabel(self.xlabel)
        plt.ylabel(self.ylabel)
        plt.grid()
        length = len(self.legend)
        for i in range(length):
            plt.plot(np.arange(len(self.data))[1:],self.data[1:,i],color = self.color[i],label = i+1)
            # plt.plot([1,2,3],[1,2,3],color = self.color[i],label = i+1)
        plt.legend(self.legend)
        plt.show()

    def eval_loss(self,loss,net,iter_data):
        losssum = 0.0
        num = 0
        for x, y in iter_data:
            y_h = net(x)
            l = loss(y_h, y)/len(y)
            losssum+=l
            num+=1
        return losssum/num
    def clear(self):
        self.data = np.zeros(shape = (1,len(self.legend)))



# my = show([1,2,3],"x","y")
# plt.ion()
# for i in range(10):
#     a = [1,2,3+i]
#     my.add(args=a)
#     print(np.arange(len(my.data))[1:])
#     print(my.data[1:,2])
#     my.show_train()
#     plt.pause(0.5)